# -*- coding: utf-8 -*-
import numpy as np
import scipy.linalg as la


def csp_train(train_x, train_y, m, filter_bank=False):
    """
    训练CSP模型，生成投影矩阵F

    输入参数
    ----------
    train_x: T×N×L ndarray
            T: 采样点数  N: 通道数  L: 训练数据trial总数
    train_y: 1 维 L 个
            L个trial对应的标签（二类）
         m: int 提取的CSP特征对数
    返回值
    ----------
    csp_proj_matrix: 2m×N
            processor 投影矩阵
    """
    channel_num = train_x.shape[1]
    if filter_bank:
        filter_num = train_x.shape[2]
        csp_proj_matrix = np.zeros([2 * m, channel_num, filter_num])
        for k in range(filter_num):
            R1, R2 = cov_matrix(train_x, train_y, k)
            [w0, w1] = la.eig(R1, R1 + R2)
            I = np.argsort(-np.real(w0))
            I = (np.hstack((I[0:m], I[channel_num - m:channel_num]))).tolist()
            csp_proj_matrix[:, :, k] = np.transpose(w1[:, I])
    else:
        R1, R2 = cov_matrix(train_x, train_y)
        [w0, w1] = la.eig(R1, R1 + R2)
        I = np.argsort(-np.real(w0))
        I = (np.hstack((I[0:m], I[channel_num - m:channel_num]))).tolist()
        csp_proj_matrix = np.transpose(w1[:, I])
    return csp_proj_matrix


def cov_matrix(train_x, train_y, filter_idx=None):
    channel_num = train_x.shape[1]
    if filter_idx is None:
        train_size = train_x.shape[2]
    else:
        train_size = train_x.shape[3]
    # 初始化R1，R2
    R1 = np.zeros([channel_num, channel_num])  # R1/R2: N×N
    R2 = np.zeros([channel_num, channel_num])
    count_l, count_r = 0, 0
    # 按数据标签分类，分别赋给R1,R2, 类1 train_y=1, 类2 train_y=0
    for i in range(train_size):
        if filter_idx is None:
            x = train_x[:, :, i]
        else:
            x = train_x[:, :, filter_idx, i]
        R = np.dot(np.transpose(x), x) / np.trace(np.dot(np.transpose(x), x))
        if train_y[i] == 1:
            R1 = R1 + R  # 求和
            count_l = count_l + 1
        else:
            R2 = R2 + R
            count_r = count_r + 1
    # 协方差矩阵的归一化
    R1, R2 = R1 / count_l, R2 / count_r
    return R1, R2